import jieba
import matplotlib.pyplot as plt
import pandas as pd
from wordcloud import WordCloud


def generate(sign=1):
    con = "POS" if sign else "NEG"
    csv_path = "../model/sentences_with_tags.csv"
    tag = "评论内容"
    df = pd.read_csv(csv_path)

    stopwords = set(
        ['的', '酒店', '在', '了', '是', '和', '我', '你', '他', '她', '它', '我们', '你们', '他们', '这个', '那个',
         '这些', '那些', "而且", "但是", "月", "日", "因为", "自己", "居然"])

    # 获取所有单词的列表
    sentences = df[df["tags"] == sign][tag].tolist()
    wordlist = []
    for sentence in sentences:
        words = jieba.cut(sentence)
        wordlist += words
    wordlist = list(set(wordlist))
    filtered_words = [word for word in wordlist if word not in stopwords]
    text_with_space = ' '.join(filtered_words)

    # 创建WordCloud对象，并设置相关参数
    wc = WordCloud(
        font_path='./hanyiokaij.ttf',  # 设置字体路径，确保支持中文
        background_color='white',  # 背景颜色
        max_words=50,  # 最大显示的词数
        max_font_size=100,  # 最大字体大小
        width=800,  # 图片宽度
        height=400  # 图片高度
    )

    wc.generate(text_with_space)

    plt.figure(figsize=(10, 5))
    plt.imshow(wc, interpolation='bilinear')
    plt.axis('off')  # 关闭坐标轴
    # plt.show()

    # 如果需要，也可以将词云图保存到文件中
    wc.to_file('{}_wordcloud.png'.format(con))


generate(0)
generate(1)
